/*
 * Copyright (c) 2024 <mewiteor@hotmail.com>
 * Monitor-Demo is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan
 * PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY
 * KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
 * NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 * 
 * @file: cpu usage.c
 * @brief: 每个进程的 CPU 使用率
 * 每秒显示一次 CPU 使用率。
 * 
 * @author: mewiteor
 * @email: mewiteor@hotmail.com
 * @date: 2024-05-26
 * @version: 1.0.0.0
 * @license: Mulan PSL v2
 */

#include <Windows.h>
#include <iostream>
#include <iomanip>
#include <cstdint>
#include <conio.h>
#include <TlHelp32.h>
#include <map>
#include <string>
#include <algorithm>

void EnablePriv();
DWORD GetCpuCount();

int main()
{
    EnablePriv();
    DWORD CPU_COUNT = GetCpuCount();

    OSVERSIONINFOW osinfo;
    GetVersionExW(&osinfo);
    DWORD flag = osinfo.dwMajorVersion >= 6 ? PROCESS_QUERY_LIMITED_INFORMATION : PROCESS_QUERY_INFORMATION;

    std::map<DWORD, uint64_t> processLastSysTimes;
    uint64_t sysLastTime = 0;

    while(_kbhit()) _getch();
    while(!_kbhit())
    {
        PROCESSENTRY32W pe32;
        HANDLE hProcessSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
        pe32.dwSize = sizeof(PROCESSENTRY32W);

        uint64_t allTime;
        GetSystemTimeAsFileTime((LPFILETIME)&allTime);

        if (INVALID_HANDLE_VALUE != hProcessSnap)
        {
            std::map<DWORD, std::wstring> allNames;
            std::multimap<double, DWORD> usage_pid;
            std::map<DWORD, uint64_t> processCurSysTimes;
            if(Process32FirstW(hProcessSnap, &pe32)) do
            {
                DWORD pid = pe32.th32ProcessID;
                uint64_t creationTime, exitTime, kernelTime, userTime;

                if (!pid)
                    continue;

                HANDLE hProcess = OpenProcess(flag, FALSE, pid);
                if (GetProcessTimes(hProcess, (LPFILETIME)&creationTime, (LPFILETIME)&exitTime, (LPFILETIME)&kernelTime, (LPFILETIME)&userTime))
                {
                    uint64_t systemTime = kernelTime + userTime;
                    allNames.insert(std::make_pair(pid, pe32.szExeFile));
                    double usage = -1;
                    if(sysLastTime)
                    {
                        auto it = processLastSysTimes.find(pid);
                        if (it != processLastSysTimes.end())
                            usage = (systemTime - it->second) * 100.0 / CPU_COUNT / (allTime - sysLastTime);
                    }
                    usage_pid.insert(std::make_pair(usage, pid));
                    processCurSysTimes.insert(std::make_pair(pid, systemTime));
                }

                CloseHandle(hProcess);
            } while (Process32NextW(hProcessSnap, &pe32));

            processLastSysTimes.swap(processCurSysTimes);

            system("cls");

            size_t i = 0;
            for (auto &it = usage_pid.crbegin(); it != usage_pid.crend(); ++it)
            {
                double usage = it->first;
                DWORD pid = it->second;
                std::wstring name = allNames[pid];
                WriteConsoleW(GetStdHandle(STD_OUTPUT_HANDLE), name.c_str(), name.size(), nullptr, nullptr);
                std::cout << ": " << pid << ": " << std::fixed << std::setprecision(1) << usage << '%' << std::endl;
                if(++i >= 10)
                    break;
            }
        }
        CloseHandle(hProcessSnap);
        sysLastTime = allTime;
        Sleep(1000);
    }
    while(_kbhit()) _getch();
    return 0;
}

void EnablePriv()
{
    HANDLE hToken;
    TOKEN_PRIVILEGES tkp;
    OpenProcessToken(GetCurrentProcess(),TOKEN_ADJUST_PRIVILEGES,&hToken);
    LookupPrivilegeValue( NULL,SE_DEBUG_NAME,&tkp.Privileges[0].Luid ); //修改进程权限
    tkp.PrivilegeCount=1;
    tkp.Privileges[0].Attributes=SE_PRIVILEGE_ENABLED;
    AdjustTokenPrivileges(hToken,FALSE,&tkp,sizeof tkp,NULL,NULL); //通知系统修改进程权限
    CloseHandle(hToken);
}

DWORD GetCpuCount()
{
	SYSTEM_INFO si;
	GetSystemInfo(&si);
    return si.dwNumberOfProcessors;
}